import os
import re
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import sklearn
from sklearn.model_selection import train_test_split
from sklearn import metrics
import skimage
import skimage.color
import skimage.io
import skimage.feature
import skimage.transform
from sklearn.base import BaseEstimator, TransformerMixin
from skimage.color import rgb2gray
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import GridSearchCV
import sklearn.metrics
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from mathe import rgb2gray_transform, hogtransformer

# load the data
data = pickle.load(open('./pickle_files/data_animals_head_20.pickle','rb'))
print("Load pickle data")

dataDesc = data['description']  # There are 20 classes and 2057 images are there. All the images are 80 x 80 (rgb)
X = data['data']                # Contains all images, 
y = data['target']              # Ground truth
labels = data['labels']         # Label names of all dataset

# split the data into train and test
x_train,x_test,y_train,y_test = train_test_split(X,y,test_size=0.2,stratify=y)
print("x_train.shape: {}; x_test.shape: {}; len(y_train): {}; len(y_test): {}".format(x_train.shape,x_test.shape,len(y_train),len(y_test)))

# 这个Pipeline里面没有设置参数，sklearn里面的pipeline有一个自动寻找最优参数的功能。
print("Modelling pipeline with automatic hyperparameter tunning.")
model_pipeline = Pipeline([
    ('grascale',rgb2gray_transform()),
    ('hogtransform',hogtransformer()),
    ('scale',StandardScaler()),
    ('sgd',SGDClassifier())
])

# 设置一个list，里面包含了可能的参数，比如hogtransform类中的orientations。注意，hogtransform和orientations之间有两个下划线
param_grid = [
    {
        'hogtransform__orientations' : [7,8,9,10],
        'hogtransform__pixels_per_cell' : [(7,7),(8,8),(9,9)],
        'hogtransform__cells_per_block' : [(2,2),(3,3)],
        'sgd__loss' : ['hinge','squared_hinge','perceptron'],
        'sgd__learning_rate': ['optimal'] 
    },
    {
        'hogtransform__orientations' : [7,8,9,10],
        'hogtransform__pixels_per_cell' : [(7,7),(8,8),(9,9)],
        'hogtransform__cells_per_block' : [(2,2),(3,3)],
        'sgd__learning_rate': ['adaptive'],
        'sgd__eta0' : [0.001,0.01,0.1]
    }
]

model_grid = GridSearchCV(model_pipeline,
        param_grid=param_grid,scoring='accuracy',
        n_jobs=-1,cv=5,verbose=2)

model_grid.fit(x_train,y_train)
# You will see 'Fitting 3 folds for each of 192 candidates, totalling 576 fits'.
# So, the system will try every hyper-parameter and compare the results. 
# Finally, system gives out the best one.

print("Best parameter is: ",model_grid.best_params_)
print("Best score is: ", model_grid.best_score_)

model_best = model_grid.best_estimator_

y_pred = model_best.predict(x_test)

# Model Evaluation
cr = sklearn.metrics.classification_report(y_test,y_pred,output_dict=True)
print(pd.DataFrame(cr).T)
print("Model evaluation score: ", metrics.cohen_kappa_score(y_test,y_pred))

# save the model
pickle.dump(model_best,open('./pickle_files/dsa_model_best.pickle','wb'))
print("Save the best model.")

